import numpy as np
from scipy.special import i0, i1  # 用于计算修正贝塞尔函数

class GainCalculator:
    def __init__(self, fft_size: int, min_suppression_level=0.1, max_suppression_level=1.0, 
                 over_subtraction_factor=2.0, prior_snr_floor=1e-6, posterior_snr_floor=1e-6):
        """
        WebRTC 风格的噪声抑制增益计算模块
        
        参数:
        fft_size: 频域频谱的大小 (n_bins = fft_size//2 + 1)
        min_suppression_level: 最小增益限制 (避免音乐噪声)
        max_suppression_level: 最大增益限制 (避免过度放大)
        over_subtraction_factor: 过减因子 (抑制更宽的噪声谱)
        prior_snr_floor: 先验信噪比下限 (避免除以零)
        posterior_snr_floor: 后验信噪比下限 (避免除以零)
        """
        self.n_bins = fft_size // 2 + 1
        self.min_suppression_level = min_suppression_level
        self.max_suppression_level = max_suppression_level
        self.over_subtraction_factor = over_subtraction_factor
        self.prior_snr_floor = prior_snr_floor
        self.posterior_snr_floor = posterior_snr_floor
        
        # 创建频带配置 (类似WebRTC)
        self._create_frequency_bands()
        
        # 设置每个频带的噪声抑制强度
        self.suppression_levels = self._set_band_suppression()
        
        # 先验信噪比缓存
        self.prior_snr = None
    
    def _create_frequency_bands(self):
        """创建频带分组 (类似WebRTC的分组方式)"""
        # WebRTC风格的频带划分
        self.band_boundaries = [0, 100, 200, 300, 400, 500, 600, 700, 800, 
                               900, 1000, 2000, 3000, 4000, 8000]  # Hz
        
        # 更细粒度的频率带定义
        self.fine_bands = [
            (0, 6),    # 最低频带
            (6, 12),   # 低频带
            (12, 50),  # 中低频带
            (50, 100), # 中频带
            (100, 200),# 中高频带
            (200, 300),# 高频带
            (300, 4000),# 最高频带
        ]
    
    def _set_band_suppression(self):
        """
        设置每个频带的噪声抑制强度
        WebRTC 对高频带应用更强的降噪
        """
        suppression_levels = np.ones(self.n_bins) * 1.0  # 默认中等抑制
        
        # 低频带：较弱抑制 (保护语音基音)
        for start, end in self.fine_bands:
            if end < 200:
                suppression_levels[start:end] = 0.9
        
        # 中频带：中等抑制 (主要语音区域)
        for start, end in self.fine_bands:
            if 200 <= start and end < 2000:
                suppression_levels[start:end] = 0.7
        
        # 高频带：强抑制 (噪声主导区域)
        for start, end in self.fine_bands:
            if start >= 2000:
                suppression_levels[start:end] = 0.5
        
        return suppression_levels
    
    def compute_gains(self, power_spectrum: np.ndarray, noise_psd: np.ndarray, 
                      prior_snr: np.ndarray = None) -> np.ndarray:
        """
        计算降噪增益 (主要入口函数)
        
        参数:
        power_spectrum: 当前帧的功率谱
        noise_psd: 噪声功率谱密度估计
        prior_snr: 可选先验信噪比 (如果不提供将计算)
        
        返回:
        增益系数数组 (范围在0-1)
        """
        # 1. 计算后验信噪比 (当前观察值)
        posterior_snr = power_spectrum / np.maximum(noise_psd, self.posterior_snr_floor)
        
        # 2. 计算或更新先验信噪比
        if prior_snr is None:
            prior_snr = self._compute_decision_directed_prior_snr(posterior_snr)
        
        # 3. 计算MMSE-STSA增益 (核心降噪增益)
        mmse_gains = self._compute_mmse_stsa_gain(prior_snr, posterior_snr)
        
        # 4. 应用过减因子和频带特定调整
        gain_factors = self._apply_over_subtraction(mmse_gains, prior_snr, posterior_snr)
        
        # 5. 应用增益限制和全局约束
        final_gains = self._apply_gain_constraints(gain_factors, posterior_snr)
        
        # 保存当前先验信噪比用于下一帧
        self.prior_snr = prior_snr.copy()
        
        return final_gains
    
    def _compute_decision_directed_prior_snr(self, posterior_snr: np.ndarray) -> np.ndarray:
        """
        使用决策导向方法计算先验信噪比
        
        参数:
        posterior_snr: 当前帧的后验信噪比
        
        返回:
        先验信噪比估计
        """
        if self.prior_snr is None:
            # 第一帧初始化
            return np.maximum(posterior_snr - 1, self.prior_snr_floor)
        
        # 决策导向递归更新
        # prev_gain = self.prior_snr / (1 + self.prior_snr) 是前一帧的增益估计
        prev_gain = self.prior_snr / np.maximum(1 + self.prior_snr, 1e-10)
        dd_prior_snr = 0.98 * (prev_gain ** 2) * posterior_snr + 0.02 * np.maximum(posterior_snr - 1, 0)
        
        return np.maximum(dd_prior_snr, self.prior_snr_floor)
    
    def _compute_mmse_stsa_gain(self, prior_snr: np.ndarray, posterior_snr: np.ndarray) -> np.ndarray:
        """
        计算基于MMSE-STSA准则的增益因子
        
        这是降噪算法的核心部分，计算复杂度较高
        """
        # 1. 计算参数v
        v_num = prior_snr * posterior_snr
        v_den = 1 + prior_snr
        v = v_num / np.maximum(v_den, 1e-10)
        
        # 2. 计算分母项 (避免数值溢出)
        exp_v = np.exp(-v / 2)
        
        # 3. 计算修正贝塞尔函数
        # 由于计算量较大，我们只在高v值区域计算全函数
        low_v_mask = v < 10
        high_v_mask = ~low_v_mask
        
        # 低v值区域
        low_v_gains = np.zeros(v.shape)
        low_v_gains[low_v_mask] = self._low_v_mmse_gain(v[low_v_mask], exp_v[low_v_mask])
        
        # 高v值区域
        high_v_gains = np.zeros(v.shape)
        high_v_gains[high_v_mask] = self._high_v_mmse_gain(v[high_v_mask])
        
        # 合并结果
        mmse_gains = np.where(low_v_mask, low_v_gains, high_v_gains)
        
        # 4. 平方根得到幅度增益
        return np.sqrt(mmse_gains)
    
    def _low_v_mmse_gain(self, v: np.ndarray, exp_v: np.ndarray) -> np.ndarray:
        """计算低v值区域的MMSE增益 (精确计算)"""
        # 计算修正贝塞尔函数 (一阶和零阶)
        i0_half = i0(v / 2)
        i1_half = i1(v / 2)
        
        # 组合计算项
        term1 = (1 + v) * i0_half
        term2 = v * i1_half
        
        # 完整公式
        return (np.pi / 4) * (v / v) * exp_v * (term1 + term2)  # v/v 处理v=0的情况
    
    def _high_v_mmse_gain(self, v: np.ndarray) -> np.ndarray:
        """计算高v值区域的近似MMSE增益 (避免复杂的贝塞尔函数计算)"""
        # 使用大v近似公式
        return 1 / (1 + 1 / np.maximum(v, 1e-10))
    
    def _apply_over_subtraction(self, gains: np.ndarray, prior_snr: np.ndarray, 
                               posterior_snr: np.ndarray) -> np.ndarray:
        """
        应用过减因子调整增益
        
        在高噪声区域应用更强的抑制
        """
        # 基于信噪比的动态过减调整
        # 低SNR区域应用更强的抑制
        snr_ratio = prior_snr / np.maximum(posterior_snr, 1e-10)
        dynamic_factor = np.minimum(self.over_subtraction_factor, 1 + self.over_subtraction_factor * (1 - snr_ratio))
        
        # 应用过减
        adjusted_gains = gains ** dynamic_factor
        
        # 应用频带特定调整
        band_adjusted = adjusted_gains * self.suppression_levels
        
        # 确保不会放大信号
        band_adjusted = np.minimum(band_adjusted, 1.0)
        
        return band_adjusted
    
    def _apply_gain_constraints(self, gains: np.ndarray, posterior_snr: np.ndarray) -> np.ndarray:
        """
        应用增益限制和全局约束
        
        包括最小增益、最大增益和语音保护策略
        """
        # 应用硬性上下限
        constrained_gains = np.clip(gains, self.min_suppression_level, self.max_suppression_level)
        
        # 语音保护策略：当后验SNR高时减少抑制
        speech_weight = 1 / (1 + np.exp(-0.5 * (posterior_snr - 5)))  # Sigmoid函数
        speech_protected = np.minimum(constrained_gains, 1 - 0.5 * (1 - constrained_gains) * speech_weight)
        
        # 平滑处理增益突变
        if self.prior_snr is not None:
            smoothed_gains = 0.9 * speech_protected + 0.1 * self.prior_snr
        else:
            smoothed_gains = speech_protected
        
        return smoothed_gains
    
    def reset(self):
        """重置模块状态"""
        self.prior_snr = None


# 测试代码
if __name__ == "__main__":
    # 模拟输入数据 (256点FFT)
    n_bins = 129
    print("创建增益计算器...")
    gain_calculator = GainCalculator(fft_size=256, min_suppression_level=0.1)
    
    # 模拟语音和噪声情况
    def simulate_speech_frame():
        """生成类似语音的功率谱"""
        freq = np.linspace(0, 8000, n_bins)
        # 基频区域强能量
        power = np.ones(n_bins) * 0.1
        power += 0.5 * np.exp(-(freq - 300)**2/(2 * 100**2))  # 基频峰
        power += 0.3 * np.exp(-(freq - 2000)**2/(2 * 500**2)) # 共振峰
        return power
    
    def simulate_noise_frame():
        """生成类似噪声的功率谱"""
        freq = np.linspace(0, 8000, n_bins)
        # 平坦噪声谱
        power = np.random.rand(n_bins) * 0.3
        # 增加高频噪声
        power += 0.1 * np.linspace(0, 1, n_bins)
        return power
    
    # 模拟语音帧
    print("\n模拟语音帧处理...")
    speech_power = simulate_speech_frame()
    noise_psd = simulate_noise_frame() * 0.5  # 噪声估计较低
    
    speech_gains = gain_calculator.compute_gains(speech_power, noise_psd)
    print(f"语音增益统计: min={np.min(speech_gains):.3f}, max={np.max(speech_gains):.3f}, mean={np.mean(speech_gains):.3f}")
    
    # 模拟噪声帧
    print("\n模拟噪声帧处理...")
    noise_power = simulate_noise_frame()
    noise_psd = simulate_noise_frame() * 0.8  # 噪声估计较高
    
    noise_gains = gain_calculator.compute_gains(noise_power, noise_psd)
    print(f"噪声增益统计: min={np.min(noise_gains):.3f}, max={np.max(noise_gains):.3f}, mean={np.mean(noise_gains):.3f}")
    
    # 可视化一个示例增益曲线
    import matplotlib.pyplot as plt
    
    plt.figure(figsize=(12, 6))
    plt.plot(np.linspace(0, 8000, n_bins), speech_gains, label="语音增益")
    plt.plot(np.linspace(0, 8000, n_bins), noise_gains, label="噪声增益")
    plt.title("增益曲线比较 (语音 vs 噪声)")
    plt.xlabel("频率 (Hz)")
    plt.ylabel("增益因子")
    plt.legend()
    plt.grid()
    plt.show()